"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China
"""

import commpy as cp
import numpy as np

# 定义调制表
mod_table_8QAM = [-1-1j, -1+1j, 1-1j, 1+1j, -1-np.sqrt(3), 1+np.sqrt(3), (-1-np.sqrt(3))*1j, (1+np.sqrt(3))*1j]

look_up_table_decoded_8QAM = [[1, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1], [0, 0, 1]]

mod_table_32QAM = [5 + 3j, 5 + 1j, 5 - 1j, 5 - 3j, 3 + 5j, 3 + 3j, 3 + 1j, 3 - 1j, 3 - 3j, 3 - 5j, 1 + 5j,
     1 + 3j, 1 + 1j, 1 - 1j, 1 - 3j, 1 - 5j, -1 + 5j, -1 + 3j, -1 + 1j, -1 - 1j, -1 - 3j, -1 - 5j,
     -3 + 5j, -3 + 3j, -3 + 1j, -3 - 1j, -3 - 3j, -3 - 5j, -5 + 3j, -5 + 1j, -5 - 1j, -5 - 3j]

look_up_table_decoded_32QAM = [[1, 0, 1, 1, 0], [1, 0, 1, 1, 1], [1, 0, 1, 0, 1], [1, 0, 1, 0, 0],
     [1, 0, 0, 1, 0], [1, 1, 1, 1, 0], [1, 1, 1, 1, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 0, 0],
     [1, 0, 0, 1, 1], [1, 1, 0, 1, 0], [1, 1, 0, 1, 1], [1, 1, 0, 0, 1], [1, 1, 0, 0, 0], [1, 0, 0, 0, 1],
     [0, 0, 0, 1, 1], [0, 1, 0, 1, 0], [0, 1, 0, 1, 1], [0, 1, 0, 0, 1], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1],
     [0, 0, 0, 1, 0], [0, 1, 1, 1, 0], [0, 1, 1, 1, 1], [0, 1, 1, 0, 1], [0, 1, 1, 0, 0], [0, 0, 0, 0, 0],
     [0, 0, 1, 1, 0], [0, 0, 1, 1, 1], [0, 0, 1, 0, 1], [0, 0, 1, 0, 0]]


def QAM_8(signal_bits, bits_per_symbol):
    X = signal_bits[0].reshape(-1, bits_per_symbol)
    data_x = []
    for i in range(len(X)):
        for j in range(2**bits_per_symbol):
            if np.array_equal(X[i, :], np.array(look_up_table_decoded_8QAM[j])):
                # data_x.append(mod_table_8QAM[j]/np.sqrt(3+np.sqrt(3)))
                data_x.append(mod_table_8QAM[j])
    data_x = np.array(data_x)

    Y = signal_bits[1].reshape(-1, bits_per_symbol)
    data_y = []
    for i in range(len(Y)):
        for j in range(2**bits_per_symbol):
            if np.array_equal(Y[i, :], np.array(look_up_table_decoded_8QAM[j])):
                # data_y.append(mod_table_8QAM[j]/np.sqrt(3+np.sqrt(3)))
                data_y.append(mod_table_8QAM[j])
    data_y = np.array(data_y)

    return data_x, data_y


def QAM_32(signal_bits, bits_per_symbol):
    X = signal_bits[0].reshape(-1, bits_per_symbol)
    data_x = []
    for i in range(len(X)):
        for j in range(2**bits_per_symbol):
            if np.array_equal(X[i, :], np.array(look_up_table_decoded_32QAM[j])):
                data_x.append(mod_table_32QAM[j])
    data_x = np.array(data_x)

    Y = signal_bits[1].reshape(-1, bits_per_symbol)
    data_y = []
    for i in range(len(Y)):
        for j in range(2**bits_per_symbol):
            if np.array_equal(Y[i, :], np.array(look_up_table_decoded_32QAM[j])):
                data_y.append(mod_table_32QAM[j])
    data_y = np.array(data_y)

    return data_x, data_y




def modulation(input, bits_per_symbol):
    """调制格式
    Args:
        input[0]: 输入X偏振bit信号, numpy类型
        input[1]: 输入Y偏振bit信号, numpy类型
        bits_per_symbol: 比特数

    Returns:
        output[0]：输出X偏振调制信号
        output[1]: 输出Y偏振调制信号
    """
    X_bits = input[0]
    Y_bits = input[1]

    if bits_per_symbol==2 or bits_per_symbol==4 or bits_per_symbol==6:
        Modulation = cp.QAMModem(2**bits_per_symbol)
        X_Pol = Modulation.modulate(X_bits)
        Y_Pol = Modulation.modulate(Y_bits)
    elif bits_per_symbol==3:
        X_Pol, Y_Pol = QAM_8(input, bits_per_symbol)
    elif bits_per_symbol == 5:
        X_Pol, Y_Pol = QAM_32(input, bits_per_symbol)

    output = [X_Pol, Y_Pol]

    return output
